import os

import pandas as pd

from app_config import get_engine_ts, get_pro

"""
valid sz50 is in the sz180
"""


def calc(date_sz50: str, date_sz180: str):
    #  000010.SH
    df_sse_50 = get_pro().index_weight(index_code='000016.SH', start_date=date_sz50, end_date=date_sz50)
    df_sse_180 = get_pro().index_weight(index_code='000010.SH', start_date=date_sz180, end_date=date_sz180)

    # 提取 df_sse_180 中的 con_code 并转换为集合
    con_codes_180 = set(df_sse_180['con_code'].unique())

    # 在 df_sse_50 中添加 isExist 列
    # 使用 lambda 函数检查每个 con_code 是否在 con_codes_180 集合中
    df_sse_50['isExist'] = df_sse_50['con_code'].apply(lambda x: x in con_codes_180)

    df_sse_50.to_excel('./file/' + date_sz50 + 'sz50.xlsx')
    df_sse_180.to_excel('./file/' + date_sz50 + 'sz180.xlsx')

    print(df_sse_50)
    print(df_sse_180)


def create_folder(folder_path):
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)
        print(f"文件夹 '{folder_path}' 创建成功")


if __name__ == '__main__':

    # calc("20250228", "20241031")
    # calc("20241031", "20240229")
    # calc("20240229", "20231031")
    # calc("20231031", "20230228")
    # calc("20230228", "20221031")
    # calc("20221031", "20220228")
    # calc("20220228", "20211029")
    # calc("20211029", "20200228")

    calc("20250228", "20250228")
    calc("20241031", "20241031")
    calc("20240229", "20240229")
    calc("20231031", "20231031")
    calc("20230228", "20230228")
    calc("20221031", "20221031")
    calc("20220228", "20220228")
    calc("20211029", "20211029")
    pass
